from scipy.integrate import solve_ivp
from itertools import product
import numpy as np
from numpy import arange
from scipy.spatial.transform import Rotation
import math

class system_model:
    def __init__(self, name):
        self.name = name

    def display_name(self):
        return f"Model Name: {self.name}"
    
    def step_interception_model(self, x0, input, run_time):
        t_range = {0, run_time}
        
        return solve_ivp(self.interception_model_formatted, t_span = t_range, y0 = x0, t_eval= arange(0, run_time, 0.001), args=input)
    
    def interception_model_formatted(self, t, x, f, omega_b):
        # set init state
        p_g = np.array(x[0:2:1]).reshape(2)
        c_z = np.array(x[2])
        v = np.array(x[3:6:1]).reshape(3,1)
        R_eb = np.array(x[6:15:1]).reshape(3,3)
        R_be = R_eb.T
        rot = Rotation.from_matrix(R_be)
        euler_rad = rot.as_euler('zyx')
        psi = euler_rad[0]
        theta = euler_rad[1]
        phi = euler_rad[2]
        

        # set parameters
        g = 9.8
        mass = 1
        dot_p_g, dot_c_z, dot_v, dot_R_eb = self.interception_model(R_eb, 
                                                                    omega_b,
                                                                    mass, 
                                                                    f, 
                                                                    g, 
                                                                    v,
                                                                    psi,
                                                                    p_g,
                                                                    c_z,
                                                                    theta,
                                                                    phi)
        dx = np.concatenate((dot_p_g.reshape(2), dot_c_z, dot_v.reshape(3), dot_R_eb.reshape(9)))
        return dx
        
    def interception_model(self, R_eb, omega_b, mass, f, g, v, psi, p_g, c_z, theta, phi):
       
        A = np.array([
            [1/c_z*math.sin(psi) + p_g[0]/c_z*math.cos(psi), -1/c_z*math.cos(psi) + p_g[0]/c_z*math.sin(psi), 0], 
            [p_g[1]/c_z*math.cos(psi), p_g[1]/c_z*math.sin(psi), -1/c_z]
            ])
        B = np.array([
            [(1+p_g[0]**2)*math.sin(theta), -(1+p_g[0]**2)*math.sin(phi)*math.sin(theta), -(1+p_g[0]**2)*math.cos(phi)*math.cos(theta)],
            [p_g[0]*p_g[1]*math.sin(theta), -p_g[0]*p_g[1]*math.sin(phi)*math.sin(theta), -p_g[0]*p_g[1]*math.cos(phi)*math.cos(theta)]
            ])
        dot_p_g = A @ v.reshape(3,1) + B @ omega_b.reshape(3,1)

        dot_c_z = -np.array([math.cos(psi), math.sin(psi), 0]) @ v.reshape(3,1)

        dot_v = np.array([[0],[0],[g]]) - f/mass * R_eb[:,2].reshape(3,1)

        dot_R_eb = R_eb @ self.vec2skew(omega_b)

        return dot_p_g, dot_c_z, dot_v, dot_R_eb
    
    def vec2skew(self, vec):
        """
        Convert a vector to a skew-symmetric matrix.
        """
        return np.array([[0, -vec[2], vec[1]],
                         [vec[2], 0, -vec[0]],
                         [-vec[1], vec[0], 0]])
    

    # def cal_D_value(self, mass:float, g:float, p_g: np.ndarray, c_z:float, v_z:float, gv_z:float, n_3:np.ndarray, v_zd:float, n_3d:np.ndarray, omega:np.ndarray, f:float) -> float:
    #     """
    #     #### Calculate the D value based on syste model.

    #     - mass: parameter, mass
    #     - g: parameter, gravity
    #     - p_g: state, 2 element np.ndarray
    #     - c_z: state, distance between the taget and the camera in the gimbal camera coordinate system 
    #     - v_z: state, vertical velocity of the body in the Earth coordinate system which is equal to the ${^g}v_y$
    #     - gv_z: state, front velocity of the body in the gimbal camera coordinate system 
    #     - n_3: state, 3 element np.ndarray, the unit vector of 
    #     - v_zd: input, desired v_z
    #     - n_3d: input, desired n_3
    #     - omega: input, desired body angular velocity in the body coordinate system
    #     - f: input, desired force in body coordinate system
    #     """
    #     A = np.array([[-1/c_z, 0, p_g[0]/c_z],
    #                   [0, -1/c_z, p_g[1]/c_z]])
    #     B = np.array([[-(1+p_g[0]**2)],
    #                   [-p_g[0]*p_g[1]]])
    #     D1 = p_g.reshape(1,2) @ A @ np.array([[0],[v_zd],[gv_z]])
    #     D2 = p_g.reshape(1,2) @ (B @ n_3.reshape(1,3) @ omega.reshape(3,1))
    #     D3 = (v_z - v_zd)*(g - f/mass*n_3[2])
    #     D4 = (n_3 - n_3d).reshape(1,3) @ (np.cross(n_3, omega)).reshape(3,1)
    #     D_value = D1 + D2 + D3 + D4 
    #     return D_value
    
    def cal_D_value(self, p_g: np.ndarray, c_z:float, gv_z:float, n_3:np.ndarray, v_zd:float, n_3d:np.ndarray, omega:np.ndarray) -> float:
        """
        #### Calculate the D value based on syste model.
        - p_g: state, 2 element np.ndarray
        - c_z: state, distance between the taget and the camera in the gimbal camera coordinate system 
        - gv_z: state, front velocity of the body in the gimbal camera coordinate system 
        - n_3: state, 3 element np.ndarray, the unit vector of 
        - v_zd: input, desired v_z
        - n_3d: input, desired n_3
        - omega: input, desired body angular velocity in the body coordinate system
        """
        A = np.array([[-1/c_z, 0, p_g[0]/c_z],
                      [0, -1/c_z, p_g[1]/c_z]])
        B = np.array([[-(1+p_g[0]**2)],
                      [-p_g[0]*p_g[1]]])
        D1 = p_g.reshape(1,2) @ A @ np.array([[0],[v_zd],[gv_z]])
        D2 = p_g.reshape(1,2) @ (B @ n_3.reshape(1,3) @ omega.reshape(3,1))
        D3 = (n_3 - n_3d).reshape(1,3) @ (np.cross(n_3, omega)).reshape(3,1)
        D_value = D1 + D2 + D3
        return D_value, D1, D2, D3
    
    def cal_D_value_n3(self, n_3:np.ndarray, omega:np.ndarray, n_3d:np.ndarray):
        """
        cal D value only for n_3→n_3d
        """
        D3 = (n_3 - n_3d).reshape(1,3) @ (np.cross(n_3, omega)).reshape(3,1)
        D_value = D3
        return D_value
    
    def cal_D_value_pg(self, p_g: np.ndarray, c_z:float, gv_z:float, n_3:np.ndarray, v_zd:float, omega:np.ndarray):
        """
        #### Calculate the D value based on syste model.
        - p_g: state, 2 element np.ndarray
        - c_z: state, distance between the taget and the camera in the gimbal camera coordinate system 
        - gv_z: state, front velocity of the body in the gimbal camera coordinate system 
        - n_3: state, 3 element np.ndarray, the unit vector of 
        - v_zd: input, desired v_z
        - omega: input, desired body angular velocity in the body coordinate system
        """
        A = np.array([[-1/c_z, 0, p_g[0]/c_z],
                      [0, -1/c_z, p_g[1]/c_z]])
        B = np.array([[-(1+p_g[0]**2)],
                      [-p_g[0]*p_g[1]]])
        D1 = p_g.reshape(1,2) @ A @ np.array([[0],[v_zd],[gv_z]])
        D2 = p_g.reshape(1,2) @ (B @ n_3.reshape(1,3) @ omega.reshape(3,1))
        D_value = D1 + D2
        return D_value, D1, D2
    

    def generate_n3(self, bound = 0.1):
        """
        Generate n3 unit vectors.
        建议生成范围在-0.2到0.2之间，步长为0.1
        """
        n3x = np.linspace(-bound, bound, num=int((bound - (-bound)) / 0.01) + 1, dtype=np.float64)
        n3y = np.linspace(-bound, bound, num=int((bound - (-bound)) / 0.01) + 1, dtype=np.float64)

        for i, j in product(range(len(n3x)), range(len(n3y))):
            n3z = math.sqrt(1 - n3x[i]**2 - n3y[j]**2)
            if not np.isnan(n3z):
                yield np.array([n3x[i], n3y[j], n3z.__round__(8)])
